from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np


# 示例文本数据
documents = [
    "向量数据库用于高维向量检索",
    "传统数据库主要用于结构化数据管理",
    "嵌入向量捕捉文本语义信息",
    "人工智能推动了向量数据库的发展"
]

# 加载预训练的Sentence-BERT模型
model = SentenceTransformer('all-MiniLM-L6-v2')
# 生成嵌入向量
embeddings = model.encode(documents)
# 打印生成的嵌入向量
print("生成的嵌入向量：")
for i, embedding in enumerate(embeddings):
    print(f"文档{i + 1}的向量：", embedding[:5], "...")  # 显示前5个维度，避免输出过长

# 计算嵌入向量之间的余弦相似度
similarity_matrix = cosine_similarity(embeddings)
# 打印相似度矩阵
print("\n嵌入向量之间的余弦相似度矩阵：")
print(similarity_matrix)

# 示例：查找与某文档最相似的文档
query_index = 0  # 假设查询第一个文档
similarities = similarity_matrix[query_index]
most_similar_index = np.argsort(similarities)[-2]  # 排除自身，取次高相似度
print(f"\n与文档{query_index + 1}最相似的文档是文档{most_similar_index + 1}，相似度为：{similarities[most_similar_index]:.4f}")
